#include <stdio.h>
#include <stdlib.h>
#include "rbtree.h"

void set_color_black(rbtree_node_t *node)
{
    node->color = COLOR_BLACK;
}

void set_color_red(rbtree_node_t *node)
{
    node->color = COLOR_RED;
}

int is_color_black(rbtree_node_t *node)
{
    if (node == NULL)
        return 1;
    return node->color == COLOR_BLACK ? 1:0;
}

int is_color_red(rbtree_node_t *node)
{
    if (node == NULL)
        return 0;
    return node->color == COLOR_RED ?1:0;
}

int is_left_child(rbtree_node_t *c)
{
    if (c->parent != NULL)
    {
        return c == c->parent->lchild?1:0;
    }
    return 0;
}

int is_right_child(rbtree_node_t *c)
{
    if (c->parent != NULL)
    {
        return c == c->parent->rchild?1:0;
    }
    return 0;
}

int is_leaf_node(rtree_node_t *n)
{
    if (n->lchild == NULL && n->rchild == NULL)
        return 1;
    return 0;
}

rbtree_node_t *new_rbtree_node(int data)
{
    rbtree_node_t *node = (rbtree_node_t *)malloc(sizeof(rbtree_node_t));
    node->data = data;
    node->parent = NULL;
    node->lchild = node->rchild = NULL;
    set_color_red(node);
}

void rbtree_node_free(rbtree_node_t *n)
{
    if(n->parent != NULL)
        n->parent = NULL;
    if(n->lchild != NULL)
        n->lchild = NULL;
    if(n->rchild != NULL)
        n->rchild = NULL;
    free(n);
}

rbtree_node_t *sibling(rbtree_node_t *c)
{
    if (c->parent == NULL)
        return NULL;
    if (c == c->parent->lchild) 
        return c->parent->rchild;
    else 
        return c->parent->lchild;
}

void left_rotate(rbtree_t *base)
{
    rbtree_node_t *s = (*base)->rchild;
    s->parent = (*base)->parent;
    if ((*base)->parent != NULL)
    {
        if (is_left_child(*base)) 
            (*base)->parent->lchild = s;
        else 
            (*base)->parent->rchild = s;
    }
    
    (*base)->parent = s;
    s->lchild = *base;
    (*base)->rchild = NULL;
    *base = s;
}

void right_rotate(rbtree_t *base)
{
    rbtree_node_t *s = (*base)->lchild;
    s->parent = (*base)->parent;
    if ((*base)->parent != NULL)
    {
        if (is_left_child(*base))
            (*base)->parent->lchild = s;
        else 
            (*base)->parent->rchild = s;
    }
   
    (*base)->parent = s;
    s->rchild = (*base);
    (*base)->lchild = NULL;
    *base = s;
}

void rebalance(rbtree_t *base)
{
    if (is_color_black(*base))
        return ;
    rbtree_node_t *parent = (*base)->parent;
    if (parent == NULL)
        return ;
    rbtree_node_t *grandpa = parent->parent;
    if (grandpa == NULL)
        return ;
    rbtree_node_t *uncle = sibling(parent);
    if (is_color_red(parent))
    {
        if (is_color_red(uncle))
        {
            set_color_black(parent);
            set_color_black(uncle);
            set_color_red(grandpa);
            *base = grandpa;
            rebalance(base);
        } 
        else 
        {
            // uncle node is black
            if (is_left_child(*base))
            {
                if(is_left_child(parent))
                {
                    set_color_black(parent);
                    set_color_red(grandpa);
                    right_rotate(&grandpa);
                    *base = grandpa;
                    return ;
                }
                else
                {
                    right_rotate(&parent);
                    *base = parent->rchild;
                    rebalance(base);
                    return ;
                }
            }
            if (is_right_child(*base))
            {
                if(is_right_child(parent))
                {
                    set_color_black(parent);
                    set_color_red(grandpa);
                    left_rotate(&grandpa);
                    *base = grandpa;
                    return ;
                }
                else
                {
                    left_rotate(&parent);
                    *base = parent->lchild;
                    rebalance(base);
                    return ;
                }
            }
        }
    }
}

rbtree_node_t *rbtree_get_left_max(rbtree_node_t *c)
{
    rbtree_node_t *t = c->lchild;
    while(t != NULL)
    {
        if(t->rchild == NULL)
            break;
        t = t->rchild;
    }
    return t;
}

int rbtree_insert(rbtree_t *root, int data)
{
    rbtree_node_t *p = new_rbtree_node(data);
    if (*root == NULL)
    {
        set_color_black(p);
        *root = p;
        return 0;
    }

    rbtree_node_t *t = *root;
    while(t != NULL)
    {
        if (data < t->data)
        {
            if (t->lchild == NULL)
            {
                t->lchild = p;
                p->parent = t;
                break;
            }
            t = t->lchild;
        }
        else if (data > t->data)
        {
            if (t->rchild == NULL)
            {
                t->rchild = p;
                p->parent = t;
                break;
            }
            t = t->rchild;
        }
        else 
        {
            printf("data %d exist.\n", data);
            return 1;
        }
    }
   
    rebalance(&p);
    if(p->parent == NULL)
    {
        set_color_black(p);
        *root = p;
    } 
    return 0;
}

rbtree_node_t *get_most_left(rbtree_node_t *n)
{
    while (n->lchild != NULL)
        n = n->lchild;
    return n;
}

// successor node
rbtree_node_t *get_successor_node(rbtree_node_t *n)
{
    if(n == NULL)
        return NULL;
    if (n->rchild != NULL)
        return get_most_left(n->rchild);
    rbtree_node_t *p = n->parent;
    while(p != NULL && p->lchild != n)
    {
        n = p;
        p = n->parent;
    }
    return p;
}

int rbtree_delete(rbtree_t *root, int data)
{
    rbtree_node_t *t = *root;
    while (t != NULL)
    {
        if (t->data == data)
        {
            break;
        } else if (data < t->data)
        {
            t = t->lchild;
        } else 
        {
            t = t->rchild;
        }
    }
    if (t == NULL)
    {
        printf("data %d not found\n", data);
        return -1;
    }
    // TODO 
    if (is_leaf_node(t))
    {
        if (is_color_red(t))
        {
            // just delete 
            if (is_left_child(t))
                t->parent->lchild = NULL;
            else 
                t->parent->rchild = NULL;
            rbtree_node_free(t);
            return 0;
        } 
        else // black 
        {
            // must fix it
            fix_rebalance(&t);
            // TODO 
            return 0;
        }
    }

    if (t->lchild != NULL)
    {
        if(t->rchild == NULL)
        {
            int tmp = t->data;
            t->data = t->lchild->data;
            t->lchild->data = tmp;
            return rbtree_delete(root, data);
        }
        // has two child 
        rbtree_node_t *n = get_successor_node(t);
        int tmp = t->data;
        t->data = n->data;
        n->data = tmp;
        return rbtree_delete(root, data);
    }

    // t->lchild == NULL
    if (t->rchild != NULL)
    {
        int tmp = t->data;
        t->data = t->rchild->data;
        t->rchild->data = tmp;
        return rbtree_delete(root, data);
    }
    return 0;
}

void fix_rebalance(rbtree_t *n)
{
    rbtree_node_t *s = sibling(*n);
    rbtree_node_t *p = (*n)->parent;
    if (is_color_black(s))
    {
        if(s->rchild != NULL)
        {
            s->color = p->color;
            set_color_black(p->rchild);
            set_color_black(s->rchild);
            left_rotate(&p);
            return ;
        }
        // TODO 
    }
}

void rbtree_inorder_print(rbtree_t root)
{
    if (root != NULL)
    {
        if (root->parent == NULL)
            printf("root:[%d<%c>] ", root->data, root->color);
        rbtree_inorder_print(root->lchild);
        printf("%d<%c>, ", root->data, root->color);
        rbtree_inorder_print(root->rchild);
    }
}

void rbtree_preorder_print(rbtree_t root)
{
    if(root != NULL)
    {
        if (root->parent == NULL)
            printf("root:[%d<%c>] ", root->data, root->color);
        printf("%d<%c>, ", root->data, root->color);
        rbtree_preorder_print(root->lchild);
        rbtree_preorder_print(root->rchild);
    }
}

void rbtree_postorder_print(rbtree_t root)
{
    if(root != NULL)
    {
        if (root->parent == NULL)
            printf("root:[%d<%c>] ", root->data, root->color);
        rbtree_postorder_print(root->lchild);
        rbtree_postorder_print(root->rchild);
        printf("%d<%c>, ", root->data, root->color);
    }
}
